/********************************************************************
	created:	30:9:2007   21:39
	author:		tiamo
	purpose:	pnp
*********************************************************************/

#include "stdafx.h"

#ifdef ALLOC_PRAGMA
	#pragma alloc_text(PAGE,AppleKeyboardStartDevice)
	#pragma alloc_text(PAGE,AppleKeyboardPnp)
#endif

// start device complete
NTSTATUS AppleKeyboardStartDeviceComplete(__in PDEVICE_OBJECT pDeviceObject,__in PIRP pIrp,__in PVOID pContext)
{
	if(pIrp->PendingReturned)
		IoMarkIrpPending(pIrp);

	PIO_STACK_LOCATION pIrpSp = IoGetCurrentIrpStackLocation(pIrp);
	if(pIrpSp->MajorFunction != IRP_MJ_PNP && pIrpSp->MinorFunction != IRP_MN_START_DEVICE)
		return STATUS_SUCCESS;

	KeSetEvent(&static_cast<CDeviceExtension*>(pContext)->m_eventStartComplete,IO_NO_INCREMENT,FALSE);
	return STATUS_MORE_PROCESSING_REQUIRED;
}

// start device
NTSTATUS AppleKeyboardStartDevice(__in CDeviceExtension* pDevExt)
{
	PAGED_CODE();

	NTSTATUS statusRet = STATUS_SUCCESS;
	PVOID pPreparsedDataBuffer = 0;
	HANDLE hKeyDevice = 0;

	__try
	{
		AppleKeyboardLoadUsageMappingList(pDevExt);

		HID_COLLECTION_INFORMATION hidCollectionInfo;
		statusRet = AppleKeyboardCallHidClass(pDevExt,IOCTL_HID_GET_COLLECTION_INFORMATION,0,0,&hidCollectionInfo,sizeof(hidCollectionInfo));
		if(!NT_SUCCESS(statusRet))
			__leave;

		AppleTraceCollectionInfo(DBGLEVEL_NONE,&hidCollectionInfo);

		pPreparsedDataBuffer = ExAllocatePoolWithTag(NonPagedPool,hidCollectionInfo.DescriptorSize,'Appl');
		if(!pPreparsedDataBuffer)
		{
			statusRet = STATUS_NO_MEMORY;
			__leave;
		}

		statusRet = AppleKeyboardCallHidClass(pDevExt,IOCTL_HID_GET_COLLECTION_DESCRIPTOR,0,0,pPreparsedDataBuffer,hidCollectionInfo.DescriptorSize);
		if(!NT_SUCCESS(statusRet))
			__leave;

		//AppleTraceCollectionData(DBGLEVEL_NONE,pPreparsedDataBuffer,hidCollectionInfo.DescriptorSize);

		PHIDP_PREPARSED_DATA pPreparsedData = static_cast<PHIDP_PREPARSED_DATA>(pPreparsedDataBuffer);
		HIDP_CAPS caps;
		statusRet = HidP_GetCaps(pPreparsedData,&caps);
		if(!NT_SUCCESS(statusRet))
			__leave;

		AppleTraceHidCaps(DBGLEVEL_NONE,&caps);

		pDevExt->m_attributeKeyboard.NumberOfKeysTotal = caps.NumberInputDataIndices;

		if(NT_SUCCESS(IoOpenDeviceRegistryKey(pDevExt->m_pPhysicalDeviceObject,PLUGPLAY_REGKEY_DEVICE,KEY_ALL_ACCESS,&hKeyDevice)))
		{
			ULONG ulValue = 0;
			if(NT_SUCCESS(AppleKeyboardQueryDeviceKey(hKeyDevice,L"KeyboardTypeOverride",&ulValue,sizeof(ulValue))))
			{
				pDevExt->m_attributeKeyboard.KeyboardIdentifier.Type = static_cast<UCHAR>(ulValue);
				pDevExt->m_wmiKeyboardIdEx.Type = ulValue;
			}

			if(NT_SUCCESS(AppleKeyboardQueryDeviceKey(hKeyDevice,L"KeyboardSubtypeOverride",&ulValue,sizeof(ulValue))))
			{
				pDevExt->m_attributeKeyboard.KeyboardIdentifier.Subtype = static_cast<UCHAR>(ulValue);
				pDevExt->m_wmiKeyboardIdEx.Subtype = ulValue;
			}

			if(NT_SUCCESS(AppleKeyboardQueryDeviceKey(hKeyDevice,L"KeyboardNumberTotalKeysOverride",&ulValue,sizeof(ulValue))))
				pDevExt->m_attributeKeyboard.NumberOfKeysTotal = static_cast<USHORT>(ulValue);

			if(NT_SUCCESS(AppleKeyboardQueryDeviceKey(hKeyDevice,L"KeyboardNumberFunctionKeysOverride",&ulValue,sizeof(ulValue))))
				pDevExt->m_attributeKeyboard.NumberOfFunctionKeys = static_cast<USHORT>(ulValue);

			if(NT_SUCCESS(AppleKeyboardQueryDeviceKey(hKeyDevice,L"KeyboardNumberIndicatorsOverride",&ulValue,sizeof(ulValue))))
				pDevExt->m_attributeKeyboard.NumberOfIndicators = static_cast<USHORT>(ulValue);

			ZwClose(hKeyDevice);
		}

		ULONG ulMaxUsageListLength = HidP_MaxUsageListLength(HidP_Input,HID_USAGE_PAGE_KEYBOARD,pPreparsedData);
		ULONG length = ulMaxUsageListLength * sizeof(USAGE_AND_PAGE) * 6 + sizeof(CHidData) + caps.InputReportByteLength;

		pDevExt->m_pHidData = static_cast<CHidData*>(ExAllocatePoolWithTag(NonPagedPool,length,'Appl'));
		if(!pDevExt->m_pHidData)
		{
			statusRet = STATUS_INSUFFICIENT_RESOURCES;
			__leave;
		}

		pDevExt->m_pHidData->m_pPreparsedData = pPreparsedData;
		pDevExt->m_pHidData->m_ulMaxUsageListLength = ulMaxUsageListLength;
		RtlCopyMemory(&pDevExt->m_pHidData->m_hidCaps,&caps,sizeof(caps));

	#define  ADD_POINTER(p,type,size) static_cast<type>(static_cast<PVOID>(static_cast<PUCHAR>(static_cast<PVOID>(p)) + size))

		pDevExt->m_pHidData->m_pInputReportBuffer = ADD_POINTER(pDevExt->m_pHidData,PVOID,sizeof(CHidData));
		pDevExt->m_pHidData->m_pPrevInputUsageList = ADD_POINTER(pDevExt->m_pHidData->m_pInputReportBuffer,PUSAGE_AND_PAGE,caps.InputReportByteLength);

		length = ulMaxUsageListLength * sizeof(USAGE_AND_PAGE);
		pDevExt->m_pHidData->m_pThisInputUsageList = ADD_POINTER(pDevExt->m_pHidData->m_pPrevInputUsageList,PUSAGE_AND_PAGE,length);
		pDevExt->m_pHidData->m_pChangedBreakUsageList = ADD_POINTER(pDevExt->m_pHidData->m_pThisInputUsageList,PUSAGE_AND_PAGE,length);
		pDevExt->m_pHidData->m_pChangedMakeUsageList = ADD_POINTER(pDevExt->m_pHidData->m_pChangedBreakUsageList,PUSAGE_AND_PAGE,length);
		pDevExt->m_pHidData->m_pMakeUsageList = ADD_POINTER(pDevExt->m_pHidData->m_pChangedMakeUsageList,PUSAGE_AND_PAGE,length);
		pDevExt->m_pHidData->m_pBreakUsageList = ADD_POINTER(pDevExt->m_pHidData->m_pMakeUsageList,PUSAGE_AND_PAGE,length);

	#undef  ADD_POINTER

		pDevExt->m_pHidData->m_mdlInputReportBuffer = IoAllocateMdl(pDevExt->m_pHidData->m_pInputReportBuffer,caps.InputReportByteLength,FALSE,FALSE,0);
		if(!pDevExt->m_pHidData->m_mdlInputReportBuffer)
		{
			statusRet = STATUS_INSUFFICIENT_RESOURCES;
			__leave;
		}

		MmBuildMdlForNonPagedPool(pDevExt->m_pHidData->m_mdlInputReportBuffer);
	}
	__finally
	{
		if(!NT_SUCCESS(statusRet))
		{
			if(pDevExt->m_pHidData)
			{
				if(pDevExt->m_pHidData->m_mdlInputReportBuffer)
					IoFreeMdl(pDevExt->m_pHidData->m_mdlInputReportBuffer);

				ExFreePool(pDevExt->m_pHidData);
			}

			if(pPreparsedDataBuffer)
				ExFreePool(pPreparsedDataBuffer);
		}
	}

	return statusRet;
}

// pnp
NTSTATUS AppleKeyboardPnp(__in PDEVICE_OBJECT pDeviceObject,__in PIRP pIrp)
{
	PAGED_CODE();

	NTSTATUS statusRet = STATUS_SUCCESS;
	CDeviceExtension* pDevExt = static_cast<CDeviceExtension*>(pDeviceObject->DeviceExtension);
	PIO_STACK_LOCATION pIrpSp = IoGetCurrentIrpStackLocation(pIrp);

	AppleDebugPrint(DBGLEVEL_INFO,"Enter AppleKeyboardPnp %p,%p\n",pDeviceObject,pIrp);

	__try
	{
		statusRet = IoAcquireRemoveLock(&pDevExt->m_removeLock,pIrp);
		if(!NT_SUCCESS(statusRet))
		{
			pIrp->IoStatus.Status = statusRet;
			IoCompleteRequest(pIrp,IO_NO_INCREMENT);
			__leave;
		}

		AppleDebugPrint(DBGLEVEL_INFO,"AppleKeyboardPnp minor code 0x%02x\n",pIrpSp->MinorFunction);
		switch(pIrpSp->MinorFunction)
		{
		case IRP_MN_START_DEVICE:
			{
				if(pDevExt->m_bDeviceStarted)
				{
					AppleDebugPrint(DBGLEVEL_INFO,"AppleKeyboard already started\n");
				}
				else
				{
					KeInitializeEvent(&pDevExt->m_eventStartComplete,NotificationEvent,FALSE);
					IoCopyCurrentIrpStackLocationToNext(pIrp);
					IoSetCompletionRoutine(pIrp,AppleKeyboardStartDeviceComplete,pDevExt,TRUE,TRUE,TRUE);
					pIrp->IoStatus.Status = STATUS_SUCCESS;
					statusRet = IoCallDriver(pDevExt->m_pLowerDevice,pIrp);
					if(statusRet == STATUS_PENDING)
						KeWaitForSingleObject(&pDevExt->m_eventStartComplete,Executive,KernelMode,FALSE,0);

					if(NT_SUCCESS(statusRet) && NT_SUCCESS(statusRet = pIrp->IoStatus.Status))
					{
						if(!pDevExt->m_bLowerKeyboardDeviceStarted)
						{
							statusRet = AppleKeyboardStartDevice(pDevExt);
							if(NT_SUCCESS(statusRet))
							{
								IoWMIRegistrationControl(pDeviceObject,WMIREG_ACTION_REGISTER);
								pDevExt->m_bLowerKeyboardDeviceStarted = TRUE;
							}
						}

						if(NT_SUCCESS(statusRet))
							pDevExt->m_bDeviceStarted = TRUE;
					}
				}

				pIrp->IoStatus.Information = 0;
				pIrp->IoStatus.Status = statusRet;
				IoCompleteRequest(pIrp,IO_NO_INCREMENT);
				IoReleaseRemoveLock(&pDevExt->m_removeLock,pIrp);
			}
			break;

		case IRP_MN_STOP_DEVICE:
			pDevExt->m_bDeviceStarted = FALSE;

		case IRP_MN_QUERY_STOP_DEVICE:
		case IRP_MN_CANCEL_STOP_DEVICE:
		case IRP_MN_QUERY_REMOVE_DEVICE:
		case IRP_MN_CANCEL_REMOVE_DEVICE:
		case IRP_MN_SURPRISE_REMOVAL:
			pIrp->IoStatus.Status = STATUS_SUCCESS;
			// fall

		default:
			IoSkipCurrentIrpStackLocation(pIrp);
			statusRet = IoCallDriver(pDevExt->m_pLowerDevice,pIrp);
			IoReleaseRemoveLock(&pDevExt->m_removeLock,pIrp);
			break;

		case IRP_MN_REMOVE_DEVICE:
			{
				AppleDebugPrint(DBGLEVEL_INFO,"Remove AppleKeyboard\n");
				IoWMIRegistrationControl(pDeviceObject,WMIREG_ACTION_DEREGISTER);

				KeCancelTimer(&pDevExt->m_timerRepeat);
				KeCancelTimer(&pDevExt->m_timerRecover);
				LARGE_INTEGER delayTime;
				delayTime.QuadPart = pDevExt->m_llTimeRepeatDelay.QuadPart;
				if(delayTime.QuadPart > pDevExt->m_llTimeRecover.QuadPart)
					delayTime.QuadPart = pDevExt->m_llTimeRecover.QuadPart;

				KeDelayExecutionThread(KernelMode,FALSE,&delayTime);
				pDevExt->m_bDeletePending = TRUE;
				KeWaitForSingleObject(&pDevExt->m_eventReadComplete,Executive,KernelMode,FALSE,0);
				IoCancelIrp(pDevExt->m_pReadIrp);
				IoSkipCurrentIrpStackLocation(pIrp);
				statusRet = IoCallDriver(pDevExt->m_pLowerDevice,pIrp);
				IoReleaseRemoveLockAndWait(&pDevExt->m_removeLock,pIrp);
				IoFreeIrp(pDevExt->m_pReadIrp);
				pDevExt->m_pReadIrp = 0;
 
				CHidData* pHidData = pDevExt->m_pHidData;
				if(pHidData)
				{
					if(pHidData->m_pPreparsedData)
						ExFreePool(pHidData->m_pPreparsedData);

					if(pHidData->m_mdlInputReportBuffer)
						IoFreeMdl(pHidData->m_mdlInputReportBuffer);

					ExFreePool(pHidData);
				}

				AppleKeyboardFreeUsageMappingList(pDevExt);
				IoDetachDevice(pDevExt->m_pLowerDevice);
				IoDeleteDevice(pDevExt->m_pFunctionDeviceObject);
			}
			break;

		}
	}
	__finally
	{
		
	}

	AppleDebugPrint(DBGLEVEL_INFO,"Leave AppleKeyboardPnp 0x%08x\n",statusRet);

	return statusRet;
}